package com.atwy.graph.uf;

/**
 * @Author: 小王子火
 * @Date: 2023/2/8
 *  优化union方法，小树接到大树上，调整树的高度
 *  这样，通过比较树的重量，就可以保证树的生长相对平衡，树的高度大致在logN这个数量级，极大提升执行效率。
 * 此时，find,union,connected的时间复杂度都下降为 O(logN)，即便数据规模上亿，所需时间也非常少。
 *
 * 我们能不能进一步压缩每棵树的高度，使树高始终保持为常数？
 */
public class UF implements UFInterface{
    /** 统计连通分量个数 */
    private int count;
    /** 节点 x 的父节点是 parent[x] */
    private int[] parent;
    /** 新增一个数组记录每棵树包含的节点数 */
    private int[] size;

    public UF(int n) {
        count = n;
        parent = new int[n];
        size = new int[n];
        for (int i = 0; i < n; i++) {
            parent[i] = i;
            // 初始设为1
            size[i] = 1;
        }
    }

    @Override
    public void union(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        if (rootP == rootQ)return;
        // 将两颗树合到一起，将小树接到大树上去
        if(size[rootP] > size[rootQ]){
            parent[rootQ] = rootP;
            size[rootP] = size[rootP] + size[rootQ];
        }else {
            parent[rootP] = rootQ;
            size[rootQ] = size[rootP] + size[rootQ];
        }
        count--;
    }

    @Override
    public boolean connected(int p, int q) {
        int rootP = find(p);
        int rootQ = find(q);
        return rootP == rootQ;
    }

    @Override
    public int count() {
        return this.count;
    }

    /** 返回某个节点 x 的根节点 */
    private int find(int p) {
        // 根节点的父节点是自己
        while (p != parent[p]){
            p = parent[p];
        }
        return p;
    }
}
